# 风格编码器（CNN）

import torch  # 必须添加在最顶部
import torch.nn as nn

class StyleEncoder(nn.Module):
    def __init__(self, latent_dim=256):
        super().__init__()
        self.resnet = resnet18(pretrained=True)
        self.resnet.fc = nn.Linear(512, latent_dim)
        self.cross_attn = nn.MultiheadAttention(latent_dim, 8)
        self.norm = nn.LayerNorm(latent_dim)

    def forward(self, batch, struct_feat):
        x = batch['image']  # 从字典获取图像
        if x.dim() == 3:
            x = x.unsqueeze(0)
        if x.shape[1] != 3:  # 确保RGB三通道
            x = x[:, :3] if x.shape[1] > 3 else x.repeat(1, 3//x.shape[1], 1, 1)
        
        visual_feat = self.resnet(x)
        struct_feat = struct_feat.mean(dim=1)
        return self.fc(torch.cat([visual_feat, struct_feat], dim=1))
        
        attn_out, _ = self.cross_attn(
            query=struct_feat,
            key=visual_feat,
            value=visual_feat
        )
        return self.norm(attn_out.squeeze(1))